#!/usr/bin/env python

import rclpy
from rclpy.node import Node
from std_msgs.msg import String, Bool
from moveit_commander import MoveGroupCommander, Constraints
from geometry_msgs.msg import PoseStamped, TransformStamped
import numpy as np
from scipy.spatial.transform import Rotation
from tf2_ros import TransformListener, Buffer
from moveit_commander.robot_trajectory import RobotTrajectory
from tf_transformations import quaternion_multiply, quaternion_inverse
from moveit_commander.robot_trajectory import RobotTrajectory
from moveit_commander.robot_trajectory import RobotTrajectory
from moveit_commander.robot_trajectory import RobotTrajectory
from tf_transformations import quaternion_matrix
import copy

############################# 全局变量 ###########################
obj_pose_msg = None
button_pose_msg = None
machine_cup_pose_msg = None
handover_pose_msg = None

tool_quat_compensate = None

class CoffeeTaskManager(Node):

    def __init__(self):
        super().__init__('coffee_task_manager')
        self.task_flag = 'wait_for_command'
        self.state_flag = 'stand_by'
        self.coffee_state = 0  # 0空 1非空

    def state_refresh(self, new_state):
        if new_state == 'grasp_stand_by':
            if self.coffee_state == 0:
                self.task_flag = "put_on_table"
            else:
                self.task_flag = "handover"
        elif new_state == 'put_finger_arrive':
            if self.coffee_state == 0:
                self.coffee_state = 1  #################### 暂时先没加左手按咖啡机
            elif self.coffee_state == 1:
                self.task_flag = 'finished'  ############# 放置且是满杯 则转为完成状态
        elif new_state == 'stand_by':
            self.task_flag = 'wait_for_command'
        self.get_logger().info(f"Coffee_Task_Manager`s state_flag was {self.state_flag} ")
        self.state_flag = new_state
        self.get_logger().info(f"Coffee_Task_Manager`s state_flag is {self.state_flag} now !!!")

# ROS 回调函数
# obj_pose 一直会保证有
def obj_pose_callback(self, msg):
    global obj_pose_msg
    obj_pose_msg = msg

    P_object = np.array([obj_pose_msg.pose.position.x, obj_pose_msg.pose.position.y, obj_pose_msg.pose.position.z])
    quaternion_object = [obj_pose_msg.pose.orientation.x, obj_pose_msg.pose.orientation.y,
                         obj_pose_msg.pose.orientation.z, obj_pose_msg.pose.orientation.w]
    # 将四元数转换为旋转矩阵
    R_object = quaternion_matrix(quaternion_object)[:3, :3]
    # 发布obj的tf 可视化
    self.tf_pub("base_link", "obj_link", P_object, R_object)


def tf_pub(self, parent_frame, child_frame, translation, rotation_matrix):
    # 发布 tf 变换 (此函数需要根据ROS2 tf2 实现)
    pass


# button_pose如果掉了 需要异常处理
# 保持上一有效数据
def button_pose_callback(self, msg):
    global button_pose_msg
    if msg.pose.orientation.x != -1 and msg.pose.orientation.y != -1:
        button_pose_msg = msg
    P_machine = np.array(
        [button_pose_msg.pose.position.x, button_pose_msg.pose.position.y, button_pose_msg.pose.position.z])
    quaternion_machine = [button_pose_msg.pose.orientation.x, button_pose_msg.pose.orientation.y,
                          button_pose_msg.pose.orientation.z, button_pose_msg.pose.orientation.w]
    # 将四元数转换为旋转矩阵
    R_machine = quaternion_matrix(quaternion_machine)[:3, :3]
    # 发布obj的tf 可视化
    self.tf_pub("base_link", "button_link", P_machine, R_machine)


def machine_cup_pose_callback(self, msg):
    global machine_cup_pose_msg
    if msg.pose.orientation.x != -1 and msg.pose.orientation.y != -1:
        machine_cup_pose_msg = msg
    P_machine = np.array([machine_cup_pose_msg.pose.position.x, machine_cup_pose_msg.pose.position.y,
                          machine_cup_pose_msg.pose.position.z])
    quaternion_machine = [machine_cup_pose_msg.pose.orientation.x, machine_cup_pose_msg.pose.orientation.y,
                          machine_cup_pose_msg.pose.orientation.z, machine_cup_pose_msg.pose.orientation.w]
    # 将四元数转换为旋转矩阵
    R_machine = quaternion_matrix(quaternion_machine)[:3, :3]
    # 发布obj的tf 可视化
    self.tf_pub("base_link", "machine_cup_link", P_machine, R_machine)


# handover_pose如果掉了 需要异常处理
# 保持上一有效数据
def handover_pose_callback(msg):
    global handover_pose_msg
    if msg.pose.orientation.x != -1 and msg.pose.orientation.y != -1:
        handover_pose_msg = msg
    P_handover = np.array(
        [handover_pose_msg.pose.position.x, handover_pose_msg.pose.position.y, handover_pose_msg.pose.position.z])
    quaternion_handover = [handover_pose_msg.pose.orientation.x, handover_pose_msg.pose.orientation.y,
                           handover_pose_msg.pose.orientation.z, handover_pose_msg.pose.orientation.w]
    # 将四元数转换为旋转矩阵
    R_handover = quaternion_matrix(quaternion_handover)[:3, :3]
    # 发布obj的tf 可视化
    tf_pub("base_link", "handover_link", P_handover, R_handover)


def upper_body_callback(msg):
    global coffee_task_manager
    if coffee_task_manager.task_flag == "wait_for_command":
        # 归位
        move_to_stand_by_pose()
        # 更新任务状态 开始任务
        coffee_task_manager.task_flag = msg.data
        task_run()
    else:
        rospy.loginfo(f"Failed to set task flag as task_flag is {coffee_task_manager.task_flag} now !!!")


def left_arm_move_to_pos(position, link_name):
    global move_group_left_arm

    current_pose = move_group_left_arm.get_current_pose().pose

    move_pose_msg = Pose()
    move_pose_msg.position.x = position[0]
    move_pose_msg.position.y = position[1]
    move_pose_msg.position.z = position[2]
    move_pose_msg.orientation.x = -0.10711095938856535  # button_pose.pose.orientation.x
    move_pose_msg.orientation.y = -0.33227448354515043  # button_pose.pose.orientation.y
    move_pose_msg.orientation.z = -0.845326815249489  # button_pose.pose.orientation.z
    move_pose_msg.orientation.w = 0.40440510059075774  # button_pose.pose.orientation.w

    print("adjusted_right_pose", move_pose_msg)
    print(current_pose)

    # 设定一个阈值，如果目标位姿与当前位置的距离大于这个值，则进行分步规划
    step_size = 0.01  # 每一步的距离，单位米

    num_steps = 4  # 计算需要的步数

    # 按照分步的策略逐步移动
    for i in range(num_steps):
        # 计算每一步的目标位置
        interpolation_factor = (i + 1) / num_steps  # 计算每一步的插值因子
        intermediate_pose = Pose()
        intermediate_pose.position.x = current_pose.position.x + (
                    position[0] - current_pose.position.x) * interpolation_factor
        intermediate_pose.position.y = current_pose.position.y + (
                    position[1] - current_pose.position.y) * interpolation_factor
        intermediate_pose.position.z = current_pose.position.z + (
                    position[2] - current_pose.position.z) * interpolation_factor

        # 使用当前的位姿插值更新目标姿态
        intermediate_pose.orientation.x = -0.10711095938856535  # button_pose.pose.orientation.x
        intermediate_pose.orientation.y = -0.33227448354515043  # button_pose.pose.orientation.y
        intermediate_pose.orientation.z = -0.845326815249489  # button_pose.pose.orientation.z
        intermediate_pose.orientation.w = 0.40440510059075774  # button_pose.pose.orientation.w

        move_group_left_arm.set_pose_target(intermediate_pose)
        plan, traj, planning_time, error_code = move_group_left_arm.plan()

        if plan:
            success = move_group_left_arm.execute(traj, wait=True)
            move_group_left_arm.stop()
            move_group_left_arm.clear_pose_targets()
            # rospy.sleep(0.2)
            if not success:
                rospy.logwarn(f"Failed to move {link_name} to intermediate pose at step {i + 1}.")
                return  # 如果某一步执行失败，则终止
        else:
            move_group_left_arm.stop()
            move_group_left_arm.clear_pose_targets()
            # rospy.sleep(0.2)
            rospy.logwarn(f"Failed to generate plan for {link_name} at step {i + 1}.")
            return  # 如果某一步规划失败，则终止

    rospy.loginfo(f"Successfully moved {link_name} to the target pose in {num_steps} steps.")


def press_button_callback():
    global move_group_left_arm
    ## 开始左臂按键任务
    button_pose = copy.deepcopy(button_pose_msg)

    button_pose_quaternion = [button_pose.pose.orientation.x, button_pose.pose.orientation.y,
                              button_pose.pose.orientation.z, button_pose.pose.orientation.w]
    button_pose_rotation_matrix = quaternion_matrix(button_pose_quaternion)

    local_offset_far = np.array(
        [-0.01, 0.01, 0.10, 1.0])  # [世界坐标系y轴（），世界坐标系z轴（数据增加z轴值减小，向物体下方移动），世界坐标系x轴（数据增加x轴值减少，离物体越远）]
    local_offset_near = np.array(
        [-0.01, 0.01, 0.00, 1.0])  # [世界坐标系y轴（），世界坐标系z轴（数据增加z轴值减小，向物体下方移动），世界坐标系x轴（数据增加x轴值减少，离物体越远）]
    world_offset_far = np.dot(button_pose_rotation_matrix, local_offset_far)
    world_offset_near = np.dot(button_pose_rotation_matrix, local_offset_near)
    # 初始位置
    pose_far = [-1, -1, -1]
    pose_far[0] = button_pose.pose.position.x + world_offset_far[0]
    pose_far[1] = button_pose.pose.position.y + world_offset_far[1]
    pose_far[2] = button_pose.pose.position.z + world_offset_far[2]
    # left_arm_move_to_pos(pose_far,"Left_arm_Link")

    rospy.sleep(0.2)

    # 按下去
    pose_naer = [-1, -1, -1]
    pose_naer[0] = button_pose.pose.position.x + world_offset_near[0]
    pose_naer[1] = button_pose.pose.position.y + world_offset_near[1]
    pose_naer[2] = button_pose.pose.position.z + world_offset_near[2]
    left_arm_move_to_pos(pose_naer, "Left_arm_Link")

    rospy.sleep(0.2)

    # 初始位置
    pose_far = [-1, -1, -1]
    pose_far[0] = button_pose.pose.position.x + world_offset_far[0]
    pose_far[1] = button_pose.pose.position.y + world_offset_far[1]
    pose_far[2] = button_pose.pose.position.z + world_offset_far[2]
    left_arm_move_to_pos(pose_far, "Left_arm_Link")


def tf_pub(frame_id, child_frame_id, P, R, is_quat=False):
    ''''''
    br = tf2_ros.TransformBroadcaster()
    t = TransformStamped()
    t.header.stamp = rospy.Time.now()
    t.header.frame_id = frame_id
    t.child_frame_id = child_frame_id

    t.transform.translation.x = P[0]
    t.transform.translation.y = P[1]
    t.transform.translation.z = P[2]

    if not is_quat:
        R_object_orientation = np.eye(4)
        R_object_orientation[:3, :3] = R
        quat = quaternion_from_matrix(R_object_orientation)
    else:
        quat = R
    t.transform.rotation.x = quat[0]
    t.transform.rotation.y = quat[1]
    t.transform.rotation.z = quat[2]
    t.transform.rotation.w = quat[3]
    br.sendTransform(t)


def compute_hand_position(P_object, R_object, dimensions, hand_y_toward):
    """
    计算夹爪接触点和姿态
    Parameters:
    - P_object: 物体位置 (3x1 array)
    - R_object: 物体姿态 (3x3 rotation matrix)
    - dimensions: 物体尺寸 (tuple), cylinder: (r, h), cuboid: (l, w, h)
    - hand_y_toward: base_link下 末端tool 的期望 y朝向（灵巧手的虎口）
    Returns:
    - P_grasp: 抓取点的位置
    - grasp_orientation: 抓取姿态 (3x3 rotation matrix)
    """

    r, h = dimensions
    obj_x_dir = R_object[:, 0]
    obj_y_dir = R_object[:, 1]
    obj_z_dir = R_object[:, 2]

    tool_x_dir = -obj_z_dir
    tool_x_dir = tool_x_dir / np.linalg.norm(tool_x_dir)

    if hand_y_toward == [0, -1, 0]:  # 'base_y_nagtive': 对应base的y轴负方向
        tool_y_candidate = np.array([0.0, -1.0, 0.0])

        dot_xy = np.abs(np.dot(tool_x_dir, tool_y_candidate))

        if dot_xy > 0.99:
            # 如果tool的x轴方向与base的y轴方向基本平行，就把tool的y轴放到base的z轴重新来一次
            tool_y_candidate = np.array([0.0, 0.0, -1.0])

        v_perp = tool_y_candidate - np.dot(tool_y_candidate, tool_x_dir) * tool_x_dir
        norm_vp = np.linalg.norm(v_perp)

        if norm_vp < 1e-8:
            tool_y_candidate = np.array([0.0, 0.0, 1.0])
            v_perp = tool_y_candidate - np.dot(tool_y_candidate, tool_x_dir) * tool_x_dir
            norm_vp = np.linalg.norm(v_perp)
    else:
        # 对应小臂朝向其他方向
        tool_y_candidate = hand_y_toward / np.linalg.norm(hand_y_toward)

        dot_xy = np.abs(np.dot(tool_x_dir, tool_y_candidate))

        if dot_xy > 0.99:
            # 如果tool的x轴方向与期望的y轴方向基本平行 说明杯子倒了。。。。
            return False, None, None
        v_perp = tool_y_candidate - np.dot(tool_y_candidate, tool_x_dir) * tool_x_dir
        norm_vp = np.linalg.norm(v_perp)

    tool_y_dir = v_perp / norm_vp
    tool_z_dir = np.cross(tool_x_dir, tool_y_dir)
    tool_z_dir = tool_z_dir / np.linalg.norm(tool_z_dir)

    grasp_orientation = np.column_stack((tool_x_dir, tool_y_dir, tool_z_dir))
    grasp_poses = []
    grasp_poses.append({
        "position": P_object + grasp_orientation @ np.array([0.03, r, 0.068]),  # 沿 y 轴正方向移动r  z轴前靠6cm
        "orientation": grasp_orientation
    })
    # 发布抓取期望位姿 tf 可视化
    tf_pub("base_link", "target_tool_link", grasp_poses[0]["position"], grasp_poses[0]["orientation"])

    return True, grasp_poses[0]["position"], grasp_poses[0]["orientation"]


def compute_new_compensation(current_pos, current_quat, target_pos, target_quat):
    """
    根据当前姿态和目标姿态计算新的补偿四元数。

    :param current_quat: [qx, qy, qz, qw]，当前姿态
    :param target_quat: [qx, qy, qz, qw]，目标姿态
    :return: [qx, qy, qz, qw]，新的补偿四元数
    """
    # 计算 delta_quat = target_quat * current_quat^-1
    current_quat_inv = quaternion_inverse(current_quat)
    delta_quat = quaternion_multiply(target_quat, current_quat_inv)

    # 计算 delta_pos = target_pos - current_pos
    delta_pos = target_pos - current_pos
    return delta_pos, delta_quat


def move_to_pose(position, quaternion, link_name, use_oc=False):
    """
    移动手臂到指定位姿。
    :param pose: Pose 消息，目标位姿
    :param link_name: 要移动的链接名称
    """
    global move_group_right_arm
    move_pose_msg = Pose()
    move_pose_msg.position.x = position[0]
    move_pose_msg.position.y = position[1]
    move_pose_msg.position.z = position[2]
    move_pose_msg.orientation.x = quaternion[0]
    move_pose_msg.orientation.y = quaternion[1]
    move_pose_msg.orientation.z = quaternion[2]
    move_pose_msg.orientation.w = quaternion[3]

    print("adjusted_right_pose", move_pose_msg)
    print(move_group_right_arm.get_current_pose().pose)

    move_group_right_arm.set_pose_target(move_pose_msg)
    plan, traj, planning_time, error_code = move_group_right_arm.plan()
    print(planning_time)

    res_execute = move_group_right_arm.execute(traj, wait=True)
    move_group_right_arm.stop()
    move_group_right_arm.clear_pose_targets()
    rospy.sleep(0.2)

    # if use_oc:
    #    move_group_right_arm.clear_path_constraints()

    if plan and res_execute:
        rospy.loginfo(f"Successfully moved {link_name} to the target pose.")
    else:
        rospy.logwarn(f"Failed to move {link_name} to the target pose.")
    return plan and res_execute


def is_pose_similar(position1, quaternion1, position2, quaternion2, max_translation_error=0.01,
                    max_rotation_error=10.0):
    """
    判断两个时刻的位姿偏差是否足够小。

    参数:
        position1 (array-like): 第一个时刻的位置 [x, y, z]。
        quaternion1 (array-like): 第一个时刻的旋转 [x, y, z, w]。
        position2 (array-like): 第二个时刻的位置 [x, y, z]。
        quaternion2 (array-like): 第二个时刻的旋转 [x, y, z, w]。
        max_translation_error (float): 平移误差阈值 (单位: 米)。
        max_rotation_error (float): 旋转误差阈值 (单位: 度)。

    返回:
        bool: 如果平移误差和旋转误差都在阈值范围内，返回 True；否则返回 False。
    """
    # 计算平移误差
    position1 = np.array(position1)
    position2 = np.array(position2)
    translation_error = np.linalg.norm(position1 - position2)

    # 计算旋转误差
    quat1 = Rotation.from_quat(quaternion1)
    quat2 = Rotation.from_quat(quaternion2)
    relative_rotation = quat1.inv() * quat2
    rotation_error = relative_rotation.magnitude() * (180 / np.pi)  # 转换为度数

    print("translation_error :  ", translation_error)
    print("rotation_error    :  ", rotation_error)

    # 判断误差是否足够小
    return translation_error <= max_translation_error and rotation_error <= max_rotation_error


def quat_multiply(quat1, quat2):
    """
    将两个四元数相乘，得到新的四元数。
    格式 qx, qy, qz, qw
    """
    qx1, qy1, qz1, qw1 = quat1
    qx2, qy2, qz2, qw2 = quat2

    # 计算四元数乘积
    qx = qw1 * qx2 + qx1 * qw2 - qy1 * qz2 + qz1 * qy2
    qy = qw1 * qy2 + qy1 * qw2 + qx1 * qz2 - qz1 * qx2
    qz = qw1 * qz2 - qx1 * qy2 + qy1 * qx2 + qz1 * qw2
    qw = qw1 * qw2 - qx1 * qx2 - qy1 * qy2 - qz1 * qz2

    return [qx, qy, qz, qw]


def move_to_stand_by_pose(is_use_compensate=False):
    # 让机械臂回归到待命位姿
    # 在任务开始、抓取完成、放置完成都会执行
    # 如果 需要补偿，则会在归位末端位姿 叠加一个 tool_quat_compensate
    # global tool_quat_compensate
    '''
    pose:       这是ARM_LINK_7的
        position:
            x: 0.4073865420857997
            y: -0.47497118260638577
            z: 0.9133771081087848
        orientation:
            x: 0.33585726536946847
            y: -0.5657318676152812
            z: 0.5082700924088744
            w: 0.5557057354623427
    '''
    # stand_by_position = [0.552,-0.271,0.961]   ## 90度 小臂
    # stand_by_quaternion = [0.6524788661440531, -0.2256188689670651, 0.6878808971115478, 0.2240252811412265] #[ 0.7071067811865475 , 0.000 , 0.7071067811865475 , 0.000 ] #虎口朝左
    stand_by_position = [0.705, -0.367, 0.951 + 0.1]
    stand_by_quaternion = [0.551, -0.388, 0.621, 0.400]

    if is_use_compensate:
        stand_by_quaternion = quat_multiply(stand_by_quaternion, tool_quat_compensate)
        # tf_pub("base_link", "tool_quat_compensate", stand_by_position, tool_quat_compensate, is_quat=True)

    # tf_pub("base_link", "stand_by_pose", stand_by_position, stand_by_quaternion, is_quat=True)

    res_execute = move_to_pose(stand_by_position, stand_by_quaternion, "Right_Arm_Link7")


def calculate_grasp_pose(P_object, R_object, dimensions, grasp_y_toward):
    '''
    计算抓取位姿
    '''
    res_cacu_grasp, P_grasp, grasp_orientation = compute_hand_position(P_object, R_object, dimensions,
                                                                       hand_y_toward=grasp_y_toward)
    if not res_cacu_grasp:
        rospy.loginfo(f"Failed to calculate grasp Pose")
    return res_cacu_grasp, P_grasp, grasp_orientation


def compute_world_offset(grasp_quaternion, pre_grasp_pose):
    '''
    根据预抓取姿态计算偏移量
    '''
    if pre_grasp_pose == 'beyond':
        local_offset = np.array([0.15, 0.0, 0.0, 1.0])
    elif pre_grasp_pose == 'front':
        local_offset = np.array([0.0, 0.1, 0.0, 1.0])
    else:
        rospy.logwarn(f"Unknown pre_grasp_pose: {pre_grasp_pose}")
        return np.zeros(4)

    return np.dot(quaternion_matrix(grasp_quaternion), local_offset)


def grasp_task(grasp_y_toward, pre_grasp_pose):
    '''
    任务主逻辑
    '''
    global coffee_task_manager
    falied_times = 0
    grasp_y_toward_copy = copy.deepcopy(grasp_y_toward)

    while coffee_task_manager.state_flag != 'grasp_hand_arrive':

        # 获取物体位姿
        P_object = np.array([obj_pose_msg.pose.position.x, obj_pose_msg.pose.position.y, obj_pose_msg.pose.position.z])
        quaternion_object = [obj_pose_msg.pose.orientation.x, obj_pose_msg.pose.orientation.y,
                             obj_pose_msg.pose.orientation.z, obj_pose_msg.pose.orientation.w]

        R_object = quaternion_matrix(quaternion_object)[:3, :3]

        # 计算抓取位置和姿态
        res_cacu_grasp, P_grasp, grasp_orientation = calculate_grasp_pose(P_object, R_object, dimensions,
                                                                          grasp_y_toward)

        if not res_cacu_grasp:
            falied_times += 1
            if falied_times >= 3:
                move_to_stand_by_pose()
                coffee_task_manager.state_refresh('stand_by')
                return False
            grasp_y_toward[0] *= 0.5
            continue

        # 计算世界偏移
        world_offset = compute_world_offset(grasp_orientation, pre_grasp_pose)

        # 构建预抓取位姿消息
        P_grasp_offset = P_grasp + world_offset[:3]

        # 执行移动
        res_execute = move_to_pose(P_grasp_offset, quaternion_from_matrix(np.eye(4)[:3, :3]), "Right_Arm_Link7",
                                   use_oc=True)
        if not res_execute:
            falied_times += 1
            if falied_times >= 3:
                move_to_stand_by_pose()
                coffee_task_manager.state_refresh('stand_by')
                return False
            continue

        falied_times = 0
        rospy.loginfo(f"Successfully moved to pre-grasp position")

        # 计算新位姿并验证是否符合要求
        new_P_object = np.array(
            [obj_pose_msg.pose.position.x, obj_pose_msg.pose.position.y, obj_pose_msg.pose.position.z])
        new_quaternion_object = [obj_pose_msg.pose.orientation.x, obj_pose_msg.pose.orientation.y,
                                 obj_pose_msg.pose.orientation.z, obj_pose_msg.pose.orientation.w]
        new_R_object = quaternion_matrix(new_quaternion_object)[:3, :3]

        res_cacu_grasp, new_P_grasp, new_grasp_orientation = calculate_grasp_pose(new_P_object, new_R_object,
                                                                                  dimensions, grasp_y_toward)

        res_similar = is_pose_similar(P_grasp, grasp_orientation, new_P_grasp, new_grasp_orientation,
                                      max_translation_error=0.05, max_rotation_error=30.0)
        if res_similar:
            coffee_task_manager.state_refresh('grasp_arm_arrive')
        else:
            continue

        res_execute = move_to_pose(new_P_grasp, quaternion_from_matrix(np.eye(4)[:3, :3]), "Right_Arm_Link7")
        if not res_execute:
            move_to_stand_by_pose()
            coffee_task_manager.state_refresh('stand_by')
            return False

        coffee_task_manager.state_refresh('grasp_hand_arrive')

    # 发布抓取信号
    msg_hand_grasp = Bool(data=False)
    hand_control_publisher.publish(msg_hand_grasp)
    rospy.loginfo(f"Publish False at {rospy.get_time()}")
    rospy.sleep(2)

    # 验证抓取并执行补偿
    ag_P_object = np.array([obj_pose_msg.pose.position.x, obj_pose_msg.pose.position.y, obj_pose_msg.pose.position.z])
    ag_quaternion_object = [obj_pose_msg.pose.orientation.x, obj_pose_msg.pose.orientation.y,
                            obj_pose_msg.pose.orientation.z, obj_pose_msg.pose.orientation.w]

    ag_R_object = quaternion_matrix(ag_quaternion_object)[:3, :3]
    res_cacu_grasp, ag_P_grasp, ag_grasp_orientation = calculate_grasp_pose(ag_P_object, ag_R_object, dimensions,
                                                                            grasp_y_toward)

    # 验算姿态误差
    res_similar = is_pose_similar(new_P_grasp, new_grasp_orientation, ag_P_grasp, ag_grasp_orientation,
                                  max_translation_error=0.05, max_rotation_error=50.0)
    if not res_similar:
        msg_hand_grasp.data = True
        hand_control_publisher.publish(msg_hand_grasp)
        rospy.loginfo(f"Publish True at {rospy.get_time()}")
        rospy.sleep(2)
        move_to_stand_by_pose()
        coffee_task_manager.state_refresh('stand_by')
        return False

    # 计算并应用补偿
    tool_pos_compensate, tool_quat_compensate = compute_new_compensation(ag_P_grasp, ag_grasp_orientation, new_P_grasp,
                                                                         new_grasp_orientation)

    move_to_stand_by_pose(is_use_compensate=True)
    coffee_task_manager.state_refresh('grasp_stand_by')

    return True


def update_pose_and_orientation(P_object, quaternion_object, dimensions, hand_y_toward):
    """更新位姿和姿态的计算"""
    R_object = quaternion_matrix(quaternion_object)[:3, :3]
    res_cacu, P_release, release_orientation = compute_hand_position(P_object, R_object, dimensions, hand_y_toward)

    release_orientation_homogeneous = np.eye(4)
    release_orientation_homogeneous[:3, :3] = release_orientation
    release_quaternion = quaternion_from_matrix(release_orientation_homogeneous)

    return P_release, release_quaternion, release_orientation

def move_to_pose_with_retry(P_target, quaternion_target, retries=3):
    """尝试移动到指定位置，失败多次则返回"""
    failed_times = 0
    while failed_times < retries:
        res_execute = move_to_pose(P_target, quaternion_target, "Right_Arm_Link7")
        if res_execute:
            return True
        failed_times += 1
    return False

def compensate_pose(P_release, release_quaternion, tool_pos_compensate, tool_quat_compensate):
    """根据补偿值更新位姿和姿态"""
    world_offset = np.dot(quaternion_matrix(release_quaternion), np.array([0.0, 0.1, 0.0, 1.0]))
    P_pre_release = P_release + world_offset[:3] + tool_pos_compensate
    pre_release_quaternion = quat_multiply(release_quaternion, tool_quat_compensate)

    return P_pre_release, pre_release_quaternion

def task_run():
    global coffee_task_manager

    # 定义物体形状和尺寸（假设是圆柱体）
    shape = 'cylinder'
    dimensions = [0.055, 0.106]  # 直径和高度

    coffee_task_manager.state_refresh('grasp_arm_go')

    res_grasp = grasp_task(grasp_y_toward=[-1, -1, 0], pre_grasp_pose='beyond')
    if not res_grasp:
        return

    falied_times = 0
    while not coffee_task_manager.state_flag == 'put_finger_arrive':
        # 获取机器位置和姿态
        P_machine_cup = np.array([machine_cup_pose_msg.pose.position.x, machine_cup_pose_msg.pose.position.y,
                                  machine_cup_pose_msg.pose.position.z])
        quaternion_machine_cup = [machine_cup_pose_msg.pose.orientation.x, machine_cup_pose_msg.pose.orientation.y,
                                  machine_cup_pose_msg.pose.orientation.z, machine_cup_pose_msg.pose.orientation.w]

        P_release, release_quaternion, release_orientation = update_pose_and_orientation(P_machine_cup,
                                                                                         quaternion_machine_cup,
                                                                                         dimensions,
                                                                                         hand_y_toward=[-1, -1, 0])

        # 计算预放置位姿
        P_pre_release, pre_release_quaternion = compensate_pose(P_release, release_quaternion, tool_pos_compensate,
                                                                tool_quat_compensate)

        # 发布预放置位姿
        tf_pub("base_link", "pre_release", P_pre_release, pre_release_quaternion, is_quat=True)

        if not move_to_pose_with_retry(P_pre_release, pre_release_quaternion):
            # 如果失败，发布手放置消息
            msg_hand_release = Bool()
            msg_hand_release.data = True
            hand_control_publisher.publish(msg_hand_release)
            rospy.sleep(2)
            return

        # 细节计算
        br_P_object = np.array(
            [obj_pose_msg.pose.position.x, obj_pose_msg.pose.position.y, obj_pose_msg.pose.position.z])
        br_quaternion_object = [obj_pose_msg.pose.orientation.x, obj_pose_msg.pose.orientation.y,
                                obj_pose_msg.pose.orientation.z, obj_pose_msg.pose.orientation.w]

        br_P_release, br_release_quaternion, br_release_orientation = update_pose_and_orientation(br_P_object,
                                                                                                  br_quaternion_object,
                                                                                                  dimensions,
                                                                                                  hand_y_toward=[-1, -1,
                                                                                                                 0])

        # 简单验算
        res_similar = is_pose_similar(P_pre_release, pre_release_quaternion, br_P_release, br_release_quaternion,
                                      max_translation_error=0.08, max_rotation_error=50.0)
        if not res_similar:
            msg_hand_grasp = Bool()
            msg_hand_grasp.data = True
            hand_control_publisher.publish(msg_hand_grasp)
            rospy.sleep(2)
            return False

        # 计算补偿
        tool_pos_compensate_inc, tool_quat_compensate_inc = compute_new_compensation(br_P_object, br_release_quaternion,
                                                                                     P_release, release_quaternion)
        tool_quat_compensate = quat_multiply(tool_quat_compensate, tool_quat_compensate_inc)
        tool_pos_compensate += tool_pos_compensate_inc

        # 更新补偿后位姿
        release_quaternion_compensate = quat_multiply(release_quaternion, tool_quat_compensate)
        P_release_compensate = P_release + tool_pos_compensate
        tf_pub("base_link", "P_release_compensate", P_release_compensate, release_quaternion_compensate, is_quat=True)

        # 手动补偿
        P_release[0] += 0.023
        P_release[1] -= 0.01
        if not move_to_pose_with_retry(P_release, release_quaternion_compensate):
            return

        coffee_task_manager.state_refresh('put_arm_arrive')

        # 发送手放置
        msg_hand_release = Bool()
        msg_hand_release.data = True
        hand_control_publisher.publish(msg_hand_release)
        rospy.sleep(2)

        # 没有反馈 直接置为到达
        coffee_task_manager.state_refresh('put_finger_arrive')

    # 后续动作
    press_button_callback()
    rospy.sleep(5)

    global move_group_right_arm
    move_group_right_arm.set_max_velocity_scaling_factor(0.3)
    move_group_right_arm.set_max_acceleration_scaling_factor(0.3)

    res_execute = move_to_pose(P_pre_release, pre_release_quaternion, "Right_Arm_Link7")
    move_to_stand_by_pose(is_use_compensate=True)
    coffee_task_manager.state_refresh('grasp_stand_by')

    # 等待handover有效到来
    while not coffee_task_manager.state_flag == 'put_finger_arrive':
        P_handover = np.array(
            [handover_pose_msg.pose.position.x, handover_pose_msg.pose.position.y, handover_pose_msg.pose.position.z])
        handover_quaternion = [handover_pose_msg.pose.orientation.x, handover_pose_msg.pose.orientation.y,
                               handover_pose_msg.pose.orientation.z, handover_pose_msg.pose.orientation.w]

        handover_quaternion_compensate = quat_multiply(handover_quaternion, tool_quat_compensate)
        P_handover_compensate = P_handover + tool_pos_compensate

        if not move_to_pose_with_retry(P_handover_compensate, handover_quaternion_compensate):
            return

        coffee_task_manager.state_refresh('put_arm_arrive')

        msg_hand_release = Bool()
        msg_hand_release.data = True
        hand_control_publisher.publish(msg_hand_release)
        rospy.sleep(2)

        coffee_task_manager.state_refresh('put_finger_arrive')

    move_to_stand_by_pose()
    coffee_task_manager.state_refresh('stand_by')


if __name__ == '__main__':
    # 初始化 ROS 节点
    rospy.init_node('coffee_task_manager_node', anonymous=True)

    # MoveIt 初始化
    move_group_right_arm = MoveGroupCommander("right_arm")
    move_group_right_arm.set_end_effector_link("right_tool_link")
    move_group_right_arm.set_planning_time(0.5)  # 0.1s 的规划用时限制
    move_group_right_arm.set_max_velocity_scaling_factor(0.2)
    move_group_right_arm.set_max_acceleration_scaling_factor(0.2)
    print(move_group_right_arm.get_planning_frame())

    # 左臂 初始化
    move_group_left_arm = MoveGroupCommander("left_arm")
    move_group_left_arm.set_end_effector_link("Left_Index_Finger_Link3")
    move_group_left_arm.set_planning_time(0.5)  # 0.1s 的规划用时限制
    move_group_left_arm.set_max_velocity_scaling_factor(0.3)
    move_group_left_arm.set_max_acceleration_scaling_factor(0.3)

    hand_control_publisher = rospy.Publisher("/hand_control_topic", Bool, queue_size=1)

    rospy.Subscriber("/obj_pose", PoseStamped, obj_pose_callback, queue_size=1)

    rospy.Subscriber("/button_pose", PoseStamped, button_pose_callback, queue_size=1)

    rospy.Subscriber("/machine_cup_pose", PoseStamped, machine_cup_pose_callback, queue_size=1)

    rospy.Subscriber("/handover_pose", PoseStamped, handover_pose_callback, queue_size=1)

    rospy.Subscriber("/upper_body_command", String, upper_body_callback, queue_size=1)

    coffee_task_manager = Coffee_Task_Manager()

    rospy.spin()

